import os
import numpy as np
from PIL import Image
from torch.utils.data import Dataset
import torch
import matplotlib.pyplot as plt
import ast

class TFs_Dataset(Dataset):
    def __init__(self, data_path, select_class=None, size=512, transform=None):
        self.data_path = data_path
        self.size = size
        self.transform = transform

        # 安全解析 select_class，如果是字符串且不是 'all'，转换为 list
        if isinstance(select_class, str):
            if select_class.lower() == 'all':
                select_class = None
            else:
                try:
                    select_class = ast.literal_eval(select_class)
                except Exception as e:
                    raise ValueError(f"select_class 解析失败，内容：{select_class}, 错误信息：{e}")
        if isinstance(select_class, list):
            select_class = [int(cls) for cls in select_class]
        self.select_class = select_class

        # 标签映射：从0开始编号
        self.label_map = {
            'T0000': 0, 'T0001': 1, 'T0010': 2, 'T0011': 3, 'T0100': 4,
            'T0101': 5, 'T0110': 6, 'T0111': 7, 'T1000': 8, 'T1001': 9,
            'T1010': 10, 'T1011': 11, 'T1100': 12, 'T1101': 13, 'T1110': 14,
            'T1111': 15, 'T10000': 16, 'T10001': 17, 'T10010': 18, 'T10011': 19,
            'T10100': 20, 'T10101': 21, 'T10110': 22, 'T10111': 23, 'T11000': 24
        }

        # 递归查找所有 .npy 文件
        self.files = []
        for root, dirs, files in os.walk(data_path):
            for file in files:
                if file.endswith('.npy'):
                    self.files.append(os.path.join(root, file))

        # 过滤符合 select_class 的样本，select_class是数字标签集合，比如[0,1,2]
        if self.select_class is not None:
            filtered_files = []
            for f in self.files:
                data = np.load(f, allow_pickle=True).item()
                label_str = data['label']
                label_num = self.label_map.get(label_str, -1)
                if label_num == -1:
                    continue  # 不在标签映射里，跳过
                if label_num in self.select_class:
                    filtered_files.append(f)
            self.files = filtered_files

        # 收集所有标签用于统计类别数
        self.labels = []
        for f in self.files:
            data = np.load(f, allow_pickle=True).item()
            label_str = data['label']
            label_num = self.label_map.get(label_str, -1)
            if label_num == -1:
                raise ValueError(f"未找到标签映射: {label_str}")
            self.labels.append(label_num)

        print(f"[TFs_Dataset] 共加载 {len(self.files)} 个样本")
        print(f"[TFs_Dataset] 共包含 {len(set(self.labels))} 个类别：{sorted(set(self.labels))}")
        # 生成 class_to_idx 属性，用于返回标签映射
        self.class_to_idx = {v: k for k, v in self.label_map.items()}

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        file_path = self.files[idx]
        data = np.load(file_path, allow_pickle=True).item()

        stft = data['stft']  # shape (2, freq, time)
        label_str = data['label']
        label_num = self.label_map.get(label_str, -1)
        if label_num == -1:
            raise ValueError(f"未找到标签映射: {label_str}")

        label_tensor = torch.tensor(label_num, dtype=torch.long)  # 标签转tensor，long类型

        # 取通道0数据 (freq, time)
        channel0 = stft[0]

        # 归一化到0~1之间，方便转图像
        channel0_norm = (channel0 - channel0.min()) / (channel0.max() - channel0.min() + 1e-6)

        # 使用jet colormap转成RGB图像
        cmap = plt.get_cmap('jet')
        channel0_rgb = cmap(channel0_norm)[:, :, :3]  # (freq, time, 3), float32 0~1

        # 转为uint8 0~255
        channel0_rgb_uint8 = (channel0_rgb * 255).astype(np.uint8)

        # 转为PIL图像
        img = Image.fromarray(channel0_rgb_uint8)

        # resize到指定尺寸
        img = img.resize((self.size, self.size))

        if self.transform is not None:
            img = self.transform(img)
        else:
            # 默认转tensor并归一化到0~1
            img = torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.

        return img, label_tensor
